import json
import pandas as pd
import numpy as np
import random
import os
import gc
import sys

seed = 123
random.seed(seed)
np.random.seed(seed)

n_fold = 5

structures = [[] for i in range(n_fold)]
targets = [[] for i in range(n_fold)]
mpids = [[] for i in range(n_fold)]

lst_len = [30000, 10000, 5000, 2348]

from pymatgen.core.structure import Structure
from collections import Counter
csv_name = './data/' + 'pbe' + '_cif.csv'
df = pd.read_csv(csv_name)
lst_len.append(len(df)-sum(lst_len))

r = list(range(len(df)))
random.shuffle(r)

cur = 0

for i in range(n_fold):
    structures[i]=df['pbe_structure'][cur:cur+lst_len[i]]
    targets[i]=df['pbe_gap'][cur:cur+lst_len[i]]
    mpids[i]=df['mp_id'][cur:cur+lst_len[i]]
    cur += lst_len[i]

print(len(structures[0]), len(targets[0]), len(mpids[0]))
print(len(structures[1]), len(targets[1]), len(mpids[1]))
print(len(structures[2]), len(targets[2]), len(mpids[2]))
print(len(structures[3]), len(targets[3]), len(mpids[3]))
print(len(structures[4]), len(targets[4]), len(mpids[4]))


# dump the data
for i in range(n_fold):
    df_pbe = pd.DataFrame({'mp_id':mpids[i], 'pbe_gap':targets[i], 'pbe_structure':structures[i]})
    df_pbe.set_index('mp_id',inplace=True)
    df_pbe.to_csv('data/'+str(i)+'_cif.csv')

